import os
import numpy as np
from sklearn.cluster import KMeans
import matplotlib.pyplot as plt

"""
Elbow Method for Optimal k
"""

def read_labels(label_file):
    objects = []
    with open(label_file, 'r') as f:
        lines = f.readlines()
    for line in lines:
        parts = line.strip().split()
        if len(parts) > 1:
            width = float(parts[3]) - float(parts[1])
            height = float(parts[4]) - float(parts[2])
            objects.append((width, height))
    return objects


def calculate_aspect_ratio(objects):
    # Avoid division by zero
    aspect_ratios = [w / h for w, h in objects if h != 0]
    return aspect_ratios


def plot_elbow_method(aspect_ratios):
    distortions = []
    K_range = range(1, 6)
    for k in K_range:
        kmeans = KMeans(n_clusters=k, random_state=0).fit(aspect_ratios)
        distortions.append(kmeans.inertia_)
    plt.plot(K_range, distortions, 'bx-',marker = 'o', markersize = 13,linewidth = 6)
    plt.xlabel('Number of Clusters (k)',fontsize = 22)

    plt.title('Elbow Method for Optimal k',fontsize = 26)
    plt.xticks(list(K_range), fontsize=24)
    plt.yticks(fontsize=24)
    plt.show()


def main():
    folder_path = '/home/hw/visdrone_labels'
    label_files = [f for f in os.listdir(folder_path) if f.endswith('.txt')]
    print("文件数量",len(label_files))

    aspect_ratios = []
    for label_file in label_files:
        objects = read_labels(os.path.join(folder_path, label_file))
        aspect_ratios.extend(calculate_aspect_ratio(objects))
    print("目标数量",len(aspect_ratios))
    aspect_ratios = np.array(aspect_ratios).reshape(-1, 1)

    # Plot elbow method
    plot_elbow_method(aspect_ratios)


if __name__ == "__main__":
    main()
